//******************************************************************************
// cmix_rsws_multifit.cc - A BASF module for reading MDST files and extracting
// quantities of interest for events consistent with a mu+mu-pi0pi0 final state.
//
// Date: June 2012
// Author: craigb (craig dot bookwalter at pnnl dot gov)
//******************************************************************************

// Standard headers
#include <algorithm>
#include <exception>
#include <iostream>
#include <fstream>
#include <map>
#include <set>
#include <stdexcept>
#include <string>
#include <sstream>
#include <utility>
#include <vector>

// Belle-specific headers
#include "particle/Particle.h"
#include "particle/PID.h"
#include "particle/utility.h"
#include "event/BelleEvent.h"
#include "basf/module.h"
#include "basf/module_descr.h"
#include "benergy/BeamEnergy.h"
#include "kid/atc_pid.h"
#include "ip/IpProfile.h"
#include "helix/Helix.h"
#include "kfitter/kvertexfitter.h"
#include "eid/eid.h"

#include <mdst/mdst.h>

#include "panther/panther.h"
#include MDST_H
#include MDST_OBS_H
#include EVTCLS_H
#include BELLETDF_H
#include RUN_INFO_H

// Headers written by me
#include "rows/cmix_countrow.hpp"
#include "rows/cmix_rswsrow.hpp"

#include <pytpp.hpp>
#include <csbbelle.hpp>
#include <csbbelle_runinforow.hpp>

using std::string;
using std::vector;
using std::cout;
using std::cerr;
using std::endl;
using std::map;

static const double RADS_TO_DEG = 360./(2*3.14159265359);
static const double D0_MASS= 1.86486;
static const double D0_SIGMA=0.0062;
static const double B_FIELD = 1.5;
static const int NFITS = 3;

#ifdef BELLE_NAMESPACE
namespace Belle {
#endif //BELLE_NAMESPACE

  class cmix_rsws_multifit: public Module {
  public:
    cmix_rsws_multifit();
    virtual ~cmix_rsws_multifit();
    void init(int *);
    void term();
    void disp_stat(const char*);
    void hist_def();
    void event(BelleEvent*, int*);
    void begin_run(BelleEvent*, int*);
    void end_run(BelleEvent*, int*);
    kvertexfitter dvtxfit(Particle& p, int& fitret, bool add_tube=false);
    void load_pmap(map<string, Particle>& pmap,
                   const Mdst_charged& slow,
                   const Mdst_charged& pos,
                   const Mdst_charged& neg);
    void fill_fields(map<string, Particle>& pmap,
                     cmix_rswsrow_t& row,
                     kvertexfitter* rskvf=NULL,
                     kvertexfitter* wskvf=NULL);

    char outpath_stub[4096];
    char inpath[4096];

    static HepPoint3D IP;
    static HepSymMatrix IPerr;
    static Ptype D0_TYPE;
    static Ptype D0B_TYPE;
    static Ptype DSTAR_TYPE;
    static Ptype DSTARB_TYPE;
    static string NT_TITLES[NFITS];
    static string NT_FILENAME_SUFFICES[NFITS];


  private:
    int _nevents_tot;
    int _combos_found;
    atc_pid  *_pid_kpi;
    Hep3Vector *_boost_cm;

    pyt::table_writer<cmix_rswsrow_t>* _nt[NFITS];
    pyt::table_writer<cmix_countrow_t>* _countnt;
    pyt::table_writer<csbbelle_runinforow_t>* _runnt;
  };

  HepPoint3D cmix_rsws_multifit::IP = HepPoint3D(0,0,0);
  HepSymMatrix cmix_rsws_multifit::IPerr = HepSymMatrix(3,0);

  Ptype cmix_rsws_multifit::D0_TYPE     = Ptype("D0");
  Ptype cmix_rsws_multifit::D0B_TYPE    = Ptype("D0B");
  Ptype cmix_rsws_multifit::DSTAR_TYPE  = Ptype("D*0");
  Ptype cmix_rsws_multifit::DSTARB_TYPE = Ptype("D*B");
  string cmix_rsws_multifit::NT_TITLES[3] =
    {"No fitting of any kind performed",
     "D0/D0B vertex fit with no tube constraint",
     "D0/D0B vertex fit with tube constraint"
    };
  string cmix_rsws_multifit::NT_FILENAME_SUFFICES[3] =
    {".no_fit.h5",
     ".dvtx_only.h5",
     ".dvtx_tube_constrained.h5"
    };

  extern "C" Module_descr *mdcl_cmix_rsws_multifit()
  {
    cmix_rsws_multifit *tmp = new cmix_rsws_multifit();
    Module_descr *dscr = new Module_descr ( "cmix_rsws_multifit", tmp);
    dscr->define_param("outpath_stub", "Path stub to output file", "S",
                       4096, tmp->outpath_stub);
    dscr->define_param("inpath", "Input file", "S", 4096, tmp->inpath);

    BeamEnergy::define_global(dscr);
    IpProfile::define_global(dscr);

    return dscr;
  }

  cmix_rsws_multifit::cmix_rsws_multifit() :
    _combos_found(0),
    _nevents_tot(0),
    _countnt(NULL),
    _runnt(NULL),
    _pid_kpi(new atc_pid(3,1,5,3,2))
  {
    for (int i=0; i < NFITS; i++)
      _nt[i] = NULL;

    // Initialize member vars
  }

  cmix_rsws_multifit::~cmix_rsws_multifit() {
    for (int i=0; i < NFITS; i++)
      delete _nt[i];

    delete _countnt;
    delete _runnt;
  }

  void cmix_rsws_multifit::init(int* status) {
    string path_stub(outpath_stub);
    string outpath;
    string title_stub = "Charm-mixing analysis with both D*+/D*- candidates";
    int expected_rows = 0;
    int complevel = 1;

    for (int i=0; i < NFITS; i++) {
      outpath = path_stub + NT_FILENAME_SUFFICES[i];
      string title = title_stub + " " + NT_TITLES[i];
      _nt[i] = new pyt::table_writer<cmix_rswsrow_t>(outpath,
                                                     "nt",
                                                     title,
                                                     expected_rows,
                                                     complevel);
      _nt[i]->table_setattr("mdst_path", inpath);
      _nt[i]->table_setattr("module", "cmix");
      _nt[i]->table_setattr("vtxfit", i);
    }

    outpath = path_stub + ".count.h5";
    string title = "Event-counting ntuple for charm-mixing analysis";
    _countnt = new pyt::table_writer<cmix_countrow_t>(outpath, "nt",
                                                      title,
                                                      expected_rows,
                                                      complevel);
    _countnt->table_setattr("mdst_path", inpath);
    _countnt->table_setattr("module", "cmix");


    outpath = path_stub + ".runinfo.h5";
    title = "Belle run information",
      _runnt = new pyt::table_writer<csbbelle_runinforow_t>(outpath, "nt", "w",
                                                            expected_rows, 0);

    cout << "CRAIGB : Analyzing file " << inpath << endl;
  }

  //--------------------------------------------------------------------------
  // hist_def() I leave empty
  //--------------------------------------------------------------------------
  void cmix_rsws_multifit::hist_def() {
    // Define ntuples here
  }


  //------------------------------------------------------------------------
  // begin_run() just contains some bookkeeping to make sure the numbers I save
  // are well-calibrated.
  //------------------------------------------------------------------------
  void cmix_rsws_multifit::begin_run(BelleEvent* be, int* status) {
    BeamEnergy::begin_run();
    IpProfile::begin_run();
    eid::init_data();

    csbbelle_runinforow_t runrow;

    runrow.expno = BeamEnergy::ExpNo();
    runrow.runno = BeamEnergy::RunNo();

    runrow.beam_e = BeamEnergy::E_beam_corr();
    runrow.e_ler = BeamEnergy::E_LER();
    runrow.e_her = BeamEnergy::E_HER();
    runrow.cross_angle = BeamEnergy::Cross_angle();

    HepLorentzVector init = initial_state();

    runrow.init_px = init.px();
    runrow.init_py = init.py();
    runrow.init_pz = init.pz();
    runrow.init_e = init.e();

    Hep3Vector boost = init.boostVector();
    runrow.boost_x = boost.x();
    runrow.boost_y = boost.y();
    runrow.boost_z = boost.z();

    _runnt->append(runrow);
  }

  void cmix_rsws_multifit::event(BelleEvent* evptr, int* status) {
    Belle_event_Manager& m_event = Belle_event_Manager::get_manager();

    _nevents_tot++;

    cmix_countrow_t countrow = cmix_countrow_default;

    countrow.expno   = m_event[0].ExpNo();
    countrow.runno   = m_event[0].RunNo();
    countrow.eventno = m_event[0].EvtNo();

    countrow.kept = 1;
    countrow.npos = -1;
    countrow.nneg = -1;

    countrow.npos_combos = -1;
    countrow.nneg_combos = -1;

    Mdst_charged_Manager& m_charged = Mdst_charged_Manager::get_manager();
    countrow.ncharged_raw = m_charged.count();

    if (IpProfile::usable()) {
      IP    = IpProfile::position(1);
      IPerr = IpProfile::position_err_b_life_smeared(1);
    }
    else {
      IP    = HepPoint3D(0,0,0);
      IPerr = HepSymMatrix(3,0);
    }

    // Build charged track combos
    vector<Mdst_charged> vpos;
    vector<Mdst_charged> vneg;

    for (Mdst_charged_Manager::iterator it = m_charged.begin();
         it != m_charged.end(); it++)
      {
        if (!good_charged(*it))
          continue;

        if (it->charge() > 0)
          vpos.push_back(*it);
        else
          vneg.push_back(*it);
      }
    countrow.npos = vpos.size();
    countrow.nneg = vneg.size();

    if ((countrow.npos + countrow.nneg) < 3) {
      countrow.kept = -1;
      _countnt->append(countrow);
      return;
    }

    HepLorentzVector initial = initial_state();
    HepLorentzVector initial_cm = initial_state();

    Hep3Vector boost_cm = initial.boostVector();
    _boost_cm      = &boost_cm;

    initial_cm.boost(-(*_boost_cm));

    vector<vector<Mdst_charged> > trios;
    track_combinations(trios, vpos, vneg, 2, 1);
    countrow.npos_combos = trios.size();
    track_combinations(trios, vpos, vneg, 1, 2);
    countrow.nneg_combos = trios.size() - countrow.npos_combos;

    if (trios.size() == 0) {
      countrow.kept = -2;
      _countnt->append(countrow);
      return;
    }

    _countnt->append(countrow);

    atc_pid pid_kpi(3,1,5,3,2);

    for (vector<vector<Mdst_charged> >::iterator it = trios.begin();
         it != trios.end(); it++)
      {
        vector<Mdst_charged> trio = *it;
        int f = 0;

        double mag_p1 = sqrt(pow(trio[0].px(), 2)
                             + pow(trio[0].py(), 2) + pow(trio[0].pz(), 2));
        double mag_p2 = sqrt(pow(trio[1].px(), 2)
                             + pow(trio[1].py(), 2) + pow(trio[1].pz(), 2));

        if (mag_p1 < mag_p2) {
          Mdst_charged tmp = trio[0];
          trio[0] = trio[1];
          trio[1] = tmp;
        }

        Mdst_charged slow = trio[1];
        Mdst_charged pos;
        Mdst_charged neg;

        map<string, Particle> pmap;
        cmix_rswsrow_t row[NFITS];
        for (int i=0; i < NFITS; i++)
          row[i] = cmix_rswsrow_default;

        if (slow.charge() > 0) {
          pos = trio[0];
          neg = trio[2];

          for (int i=0; i < NFITS; i++)
            row[i].slow_charge = 1;

        } else {
          pos = trio[2];
          neg = trio[0];

          for (int i=0; i < NFITS; i++)
            row[i].slow_charge = -1;
        }

        load_pmap(pmap, slow, pos, neg);

        double m_ws = pmap["ws"].mass();
        double m_rs = pmap["rs"].mass();

        if ( ((m_ws < 1.75) || (m_ws > 2.0)) &&
             ((m_rs < 1.75) || (m_rs > 2.0)) )
          continue;

        double q_ws = pmap["dstar_ws"].mass()
          - pmap["ws"].mass() - 0.13957018;

        double q_rs = pmap["dstar_rs"].mass()
          - pmap["rs"].mass() - 0.13957018;

        if ((q_rs > 0.02) && (q_ws > 0.02))
          continue;


        for (int i=0; i < NFITS; i++) {
          row[i].expno   = m_event[0].ExpNo();
          row[i].runno   = m_event[0].RunNo();
          row[i].eventno = m_event[0].EvtNo();

          row[i].pos_svd_nhits_rphi  = pos.trk().mhyp(2).nhits(3);
          row[i].pos_svd_nhits_z     = pos.trk().mhyp(2).nhits(4);
          row[i].pos_svd_hitarr      = pos.trk().mhyp(2).hit_svd();

          row[i].neg_svd_nhits_rphi  = neg.trk().mhyp(2).nhits(3);
          row[i].neg_svd_nhits_z     = neg.trk().mhyp(2).nhits(4);
          row[i].neg_svd_hitarr      = neg.trk().mhyp(2).hit_svd();

          row[i].slow_svd_nhits_rphi  = slow.trk().mhyp(2).nhits(3);
          row[i].slow_svd_nhits_z     = slow.trk().mhyp(2).nhits(4);
          row[i].slow_svd_hitarr      = slow.trk().mhyp(2).hit_svd();

          row[i].slow_kpi_llh = pid_kpi.prob(slow);

          row[i].slow_px = slow.px();
          row[i].slow_py = slow.py();
          row[i].slow_pz = slow.pz();

          row[i].idhash = three_body_idhash(slow.get_ID(),
                                            pos.get_ID(),
                                            neg.get_ID());

        }

        // Fill fields for the no-fit row and advance the counter.
        fill_fields(pmap, row[f++]);

        // Copy particle map so that the tubed fit is done with the
        // original values.
        map<string, Particle> pmap_tubed;
        for (map<string, Particle>::iterator it = pmap.begin();
             it != pmap.end(); it++)
          pmap_tubed[it->first] = it->second;

        kvertexfitter rskvf = dvtxfit(pmap["rs"], row[f].rs_dvf_ret);
        kvertexfitter wskvf = dvtxfit(pmap["ws"], row[f].ws_dvf_ret);
       
        pmap["rsk"] = pmap["rs"].child(0);
        pmap["rspi"] = pmap["rs"].child(1);
        pmap["wsk"] = pmap["ws"].child(0);
        pmap["wspi"] = pmap["ws"].child(1);

        // Fill fields for the vertex-fit-only row and advance the counter
        fill_fields(pmap, row[f++], &rskvf, &wskvf);

        if (IpProfile::usable()) {
          kvertexfitter rskvf_tubed =
            dvtxfit(pmap_tubed["rs"], row[f].rs_dvf_ret, true);
          kvertexfitter wskvf_tubed =
            dvtxfit(pmap_tubed["ws"], row[f].ws_dvf_ret, true);

          fill_fields(pmap_tubed, row[f++], &rskvf_tubed, &wskvf_tubed);
        }

        for (int i=0; i < f; i++)
          _nt[i]->append(row[i]);

        _combos_found++;
      }
  }


  void cmix_rsws_multifit::load_pmap(map<string, Particle>& pmap,
                                     const Mdst_charged& slow,
                                     const Mdst_charged& pos,
                                     const Mdst_charged& neg)
  {
    pmap["kp"]  = Particle(pos, "K+");
    pmap["km"]  = Particle(neg, "K-");
    pmap["pip"] = Particle(pos, "PI+");
    pmap["pim"] = Particle(neg, "PI-");

    Ptype d_type;
    Ptype dstar_type;

    if (slow.charge() > 0) {
      pmap["rsk"]  = pmap["km"];  // right-sign K
      pmap["wsk"]  = pmap["kp"];  // wrong-sign K
      pmap["rspi"] = pmap["pip"];
      pmap["wspi"] = pmap["pim"];

      pmap["slow_pi"] = Particle(slow, "PI+");
      pmap["slow_k"]  = Particle(slow, "K+");

      d_type = D0_TYPE;
      dstar_type = DSTAR_TYPE;
    } else {
      pmap["rsk"]  = pmap["kp"];  // right-sign K
      pmap["wsk"]  = pmap["km"];  // wrong-sign K
      pmap["rspi"] = pmap["pim"];
      pmap["wspi"] = pmap["pip"];

      pmap["slow_pi"] = Particle(slow, "PI-");
      pmap["slow_k"]  = Particle(slow, "K-");

      d_type = D0B_TYPE;
      dstar_type = DSTARB_TYPE;
    }

    pmap["rs"] = Particle(pmap["rsk"].p() + pmap["rspi"].p(),
                          d_type);
    pmap["rs"].relation().append(pmap["rsk"]);
    pmap["rs"].relation().append(pmap["rspi"]);

    pmap["ws"] = Particle(pmap["wsk"].p() + pmap["wspi"].p(),
                          d_type);
    pmap["ws"].relation().append(pmap["wsk"]);
    pmap["ws"].relation().append(pmap["wspi"]);

    pmap["dstar_rs"] = Particle(pmap["rs"].p() + pmap["slow_pi"].p(),
                                dstar_type);
    pmap["dstar_rs"].relation().append(pmap["rs"]);
    pmap["dstar_rs"].relation().append(pmap["slow_pi"]);

    pmap["dstar_ws"] = Particle(pmap["ws"].p() + pmap["slow_pi"].p(),
                                dstar_type);
    pmap["dstar_ws"].relation().append(pmap["ws"]);
    pmap["dstar_ws"].relation().append(pmap["slow_pi"]);
  }

  void cmix_rsws_multifit::fill_fields(map<string, Particle>& pmap,
                                       cmix_rswsrow_t& row,
                                       kvertexfitter* rskvf,
                                       kvertexfitter* wskvf)
  {
    // Fill some RS fields
    row.rspi_kpi_llh = _pid_kpi->prob(pmap["rspi"].mdstCharged());
    row.rsk_kpi_llh  = _pid_kpi->prob(pmap["rsk"].mdstCharged());

    // eid pos_eid(pos);
    // _row.pos_eid = pos_eid.prob(3, -1, 5);

    row.rspi_px = pmap["rspi"].px();
    row.rspi_py = pmap["rspi"].py();
    row.rspi_pz = pmap["rspi"].pz();

    row.rsk_px = pmap["rsk"].px();
    row.rsk_py = pmap["rsk"].py();
    row.rsk_pz = pmap["rsk"].pz();

    // Fill same fields for WS combos
    row.wspi_kpi_llh = _pid_kpi->prob(pmap["wspi"].mdstCharged());
    row.wsk_kpi_llh  = _pid_kpi->prob(pmap["wsk"].mdstCharged());

    // eid neg_eid(neg);
    // row.neg_eid = neg_eid.prob(3, -1, 5);

    row.wspi_px = pmap["wspi"].px();
    row.wspi_py = pmap["wspi"].py();
    row.wspi_pz = pmap["wspi"].pz();

    row.wsk_px = pmap["wsk"].px();
    row.wsk_py = pmap["wsk"].py();
    row.wsk_pz = pmap["wsk"].pz();


    row.m_kk   = (pmap["rsk"].p() + pmap["wsk"].p()).m();
    row.m_pipi = (pmap["rspi"].p() + pmap["wspi"].p()).m();
    row.m_rs   = pmap["rs"].mass();
    row.m_ws   = pmap["ws"].mass();

    row.m_pipipi =
      (pmap["slow_pi"].p() + pmap["rspi"].p() + pmap["wspi"].p()).m();
    row.m_kkk    =
      (pmap["slow_k"].p() + pmap["rsk"].p() + pmap["wsk"].p()).m();

    row.m_dstar_rs  = pmap["dstar_rs"].mass();
    row.m_dstar_ws  = pmap["dstar_ws"].mass();

    row.m_kkpi_rs   =
      (pmap["slow_k"].p() + pmap["rsk"].p() + pmap["rspi"].p()).m();
    row.m_kkpi_ws   =
      (pmap["slow_k"].p() + pmap["wsk"].p() + pmap["wspi"].p()).m();

    row.q_ws = pmap["dstar_ws"].mass()
      - pmap["ws"].mass() - 0.13957018;

    row.q_rs = pmap["dstar_rs"].mass()
      - pmap["rs"].mass() - 0.13957018;

    Hep3Vector rs_boost =
      (pmap["rsk"].p() + pmap["rspi"].p()).boostVector();
    HepLorentzVector p4rsk_rsrf(pmap["rsk"].px(), pmap["rsk"].py(),
                                pmap["rsk"].pz(), pmap["rsk"].e());
    HepLorentzVector p4rspi_rsrf(pmap["rspi"].px(), pmap["rspi"].py(),
                                 pmap["rspi"].pz(), pmap["rspi"].e());
    p4rsk_rsrf.boost(-rs_boost);
    p4rspi_rsrf.boost(-rs_boost);

    row.rs_theta_star = rs_boost.unit().dot(p4rspi_rsrf.vect().unit());

    Hep3Vector ws_boost =
      (pmap["wsk"].p() + pmap["wspi"].p()).boostVector();
    HepLorentzVector p4wsk_wsrf(pmap["wsk"].px(), pmap["wsk"].py(),
                                pmap["wsk"].pz(), pmap["wsk"].e());
    HepLorentzVector p4wspi_wsrf(pmap["wspi"].px(), pmap["wspi"].py(),
                                 pmap["wspi"].pz(), pmap["wspi"].e());
    p4wsk_wsrf.boost(-ws_boost);
    p4wspi_wsrf.boost(-ws_boost);

    row.ws_theta_star = ws_boost.unit().dot(p4wspi_wsrf.vect().unit());

    HepLorentzVector p4cm_dstar_rs =
      pmap["slow_pi"].p() + pmap["rsk"].p() + pmap["rspi"].p();
    HepLorentzVector p4cm_dstar_ws =
      pmap["slow_pi"].p() + pmap["wsk"].p() + pmap["wspi"].p();

    p4cm_dstar_rs.boost(-(*_boost_cm));
    p4cm_dstar_ws.boost(-(*_boost_cm));

    row.pcm_dstar_rs = p4cm_dstar_rs.vect().mag();
    row.pcm_dstar_ws = p4cm_dstar_ws.vect().mag();

    if (rskvf) {
      if (row.rs_dvf_ret == 0) {
	row.rs_dvf_cl     = rskvf->cl();
	row.rs_dvf_chi2   = rskvf->chisq();
	row.rs_dvf_ndof   = rskvf->dgf();
	
	row.rs_dvf_m_unc  = pmap["rs"].momentum().dMass();
      } 
    }
    else
      row.rs_dvf_ret = 0;
    
       

    if (wskvf) {
      if (row.ws_dvf_ret == 0) {
	row.ws_dvf_cl     = wskvf->cl();
	row.ws_dvf_chi2   = wskvf->chisq();
	row.ws_dvf_ndof   = wskvf->dgf();
	
	row.ws_dvf_m_unc  = pmap["ws"].momentum().dMass();
      }
    }
    else
      row.ws_dvf_ret = 0;

  }


  kvertexfitter cmix_rsws_multifit::dvtxfit(Particle& d, int& fitret,
                                            bool add_tube)
  {
    kvertexfitter kvf;
    kvf.magneticField(B_FIELD);
    addTrack2fit(kvf, d.relation().child(0));
    addTrack2fit(kvf, d.relation().child(1));

    if (add_tube && IpProfile::usable())
      addTube2fit(kvf);

    fitret = kvf.fit();

    if (fitret == 0) {
      kmakemother dkmm;
      dkmm.magneticField(B_FIELD);
      makeMother(dkmm, kvf, d, 0);     
    }

    return kvf;
  }


  void cmix_rsws_multifit::end_run(BelleEvent* be, int* status) {
    // Yessir
  }

  void cmix_rsws_multifit::term() {
    cout << "CRAIGB : " << _nevents_tot << " events processed, ";
    cout <<  _combos_found << " D0/D0bar candidates found" << endl;

    // End of analysis, close ntuples
  }

  void cmix_rsws_multifit::disp_stat(const char* str) {
    // Who knows
  }


#ifdef BELLE_NAMESPACE
}
#endif
